import numpy as np
from easySample import easySample
import shelve
import torch
import sys
sys.path.append(r"/home/cyw/projects/function_sim_project/basic_script")
sys.path.append(r"/home/cyw/projects/function_sim_project/function_sim")


class sample_function_embedding():
    """
        获得样本经过模型训练后的各个函数的嵌入值(新增操作每一层)
        input:模型，样本a，样本b
        output:两个样本的图节点的嵌入值
    """

    def __init__(self, model, namea, nameb) -> None:
        """
            namea为样本的MD5值
        """
        # shelve 名字区分大小写,大写的时候生成了一个新的shelve文件
        namea = namea.lower()
        nameb = nameb.lower()

        self.model = torch.load(model)
        sample = easySample()

        sample_a = sample.get_sample(namea, "functionSim")
        sample_b = sample.get_sample(nameb, "functionSim")

        x_adj = torch.tensor([sample_a["adj"]], dtype=torch.float64).cuda()
        x_att = torch.tensor([sample_a["att"]], dtype=torch.float64).cuda()

        x_type = torch.tensor([sample_a["vtype"]], dtype=torch.float64).cuda()
        y_adj = torch.tensor([sample_b["adj"]], dtype=torch.float64).cuda()
        y_att = torch.tensor([sample_b["att"]], dtype=torch.float64).cuda()
        y_type = torch.tensor([sample_b["vtype"]], dtype=torch.float64).cuda()
        score1 = self.model(x_adj, x_att, x_type, y_adj, y_att, y_type)
        print(score1)


if __name__ == "__main__":
    # 需要开启保存embedding，使用后请关闭，否则磁盘会被塞满
    # !!!!!!!!切换的时候配置信息中的设置也需要改变！！！！！！！！！！  难得写代码了，用的时候手动注释一下！！！！！！
    functionSimModel = [r"/home/cyw/projects/function_sim_project/all_data/models/functionSim_model_cross_edge_best_4_8.pth",
                        r"/home/cyw/projects/function_sim_project/all_data/models/functionSim_with_hete_model.pth",
                        r"/home/cyw/projects/function_sim_project/all_data/models/functionSim_zero_model.pth"]

    aList = [["eec55f1295d0cce40bbba7be231fdddb", "6abf24e3aa3e971e6461bcfea35fe950"],
             ["e940f6319ad362ec49af956cafaf8752",
                 "d72fe801683ee70dccaba657984b7997"],
             ["1da46f95579e467013e94ca4a5117de0",
                 "496430ce3d8623147be33937f92265a7"],
             ["7d9d0d28eaed286c029947e585993b26",
                 "1da46f95579e467013e94ca4a5117de0"],
             ["cfc34d91c9114483f3101c5049afb10d",
                 "e8a683916a98d7f132fd8d05d52ea66c"],

             ["74ec6006ed76a9bed14c140e53c7efd4",
                 "bb6f9c2aa88245746e24d7672016afb7"],
             ["b68a07cf3f63172901ce8298e78d7a74",
                 "63ba29be215b532731392511dcd4363e"],
             ["16f8a4cdd845bc4680c2822d8dbc10c4",
                 "bb6f9c2aa88245746e24d7672016afb7"],
             ["698ff16b0e824db5874acee0f7324bb2",
                 "a6bd7c5a22dbf503e7d0510b02494a95"],
             ["ec45bbfb38da5b65a58eb018b0122285",
                 "d160839802998ccabdb7ca066ae2eb3b"],

             ["50c64842dc9e99f259b4a247762b918e",
                 "bb6f9c2aa88245746e24d7672016afb7"],
             ["00a109d00d01fc2be04738ee9dcc65c2",
                 "22c1b1924703ad7d0e142c5c8ebc57e0"],
             ["81cbb6f4b1cc8d5e2167d106657de095",
                 "b33da0446151ac45b2e07491adde796d"],
             ["ce1a3ebdde8eb6d24e2c91bd3ce1da72",
                 "c1eb31c7c0ff87c5b182f8e3a447028a"],
             ["2db1afa1d21b8bbe83d7208f4c9942b3",
                 "e3fb5232f88b4a62b0d80d31e4828770"],

             ["f846e2ca4c4a36a00dbc28abd892e4b9",
                 "cbb2da21894d34cdbadaa26d1ec03c01"]
             ]
    ind = 0
    model = functionSimModel[2]

    sample_a = aList[ind][0]
    sample_b = aList[ind][1]
    a = sample_function_embedding(model, sample_a, sample_b)
    print("!!!!!!!!!!请注意functionSim目录下的配置文件是否修改!!!!!!!!!!!!!!!!")
    print("!!!!!!!!!!请注意functionSim目录下的配置文件是否修改!!!!!!!!!!!!!!!!")
    print("!!!!!!!!!!请注意functionSim目录下的配置文件是否修改!!!!!!!!!!!!!!!!")
